{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt\n",
      "1122304/1115394 [==============================] - 0s 0us/step\n",
      "length of text:1115394 characters.\n"
     ]
    }
   ],
   "source": [
    "path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')\n",
    "text=open(path_to_file,'rb').read().decode(encoding='utf-8')\n",
    "print(\"length of text:{} characters.\".format(len(text)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You are all resolved rather to die than to famish?\n",
      "\n",
      "All:\n",
      "Resolved. resolved.\n",
      "\n",
      "First Citizen:\n",
      "First, you know Caius Marcius is chief enemy to the people.\n",
      "\n",
      "All:\n",
      "We know't, we know't.\n",
      "\n",
      "First Citizen:\n",
      "Let us kill him, and we'll have corn at our own price.\n",
      "Is't a verdict?\n",
      "\n",
      "All:\n",
      "No more talking on't; let it be done: away, away!\n",
      "\n",
      "Second Citizen:\n",
      "One word, good citizens.\n",
      "\n",
      "First Citizen:\n",
      "We are accounted poor citizens, the patricians good.\n",
      "What authority surfeits on would relieve us: if they\n",
      "would yield us but the superfluity, while it were\n",
      "wholesome, we might guess they \n"
     ]
    }
   ],
   "source": [
    "print(text[:666])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unique characters:65\n"
     ]
    }
   ],
   "source": [
    "vocab=sorted(set(text))\n",
    "print('unique characters:{}'.format(len(vocab)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['\\n', ' ', '!', '$', '&', \"'\", ',', '-', '.', '3']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1115394"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 字符到数字\n",
    "char2idx={u:i for (i,u) in enumerate(vocab)}\n",
    "# 数字映射到字符\n",
    "idx2char=np.array(vocab)\n",
    "# 将文本转换为数字\n",
    "text_as_int=np.array([char2idx[c] for c in text])\n",
    "len(text_as_int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('int32')"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_as_int.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<zip at 0x156f3457e08>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# zip() 函数用于将可迭代的对象作为参数，将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的列表。\n",
    "a = [1,2,3]\n",
    "b = [4,5,6]\n",
    "c = [4,5,6,7,8]\n",
    "zipped = zip(a,b)     # 打包为元组的列表\n",
    "# [(1, 4), (2, 5), (3, 6)]\n",
    "zip(a,c)              # 元素个数与最短的列表一致\n",
    "# [(1, 4), (2, 5), (3, 6)]\n",
    "zip(*zipped)          # 与 zip 相反，*zipped 可理解为解压，返回二维矩阵式\n",
    "# [(1, 2, 3), (4, 5, 6)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "'\\n':   0,\n",
      "' ' :   1,\n",
      "'!' :   2,\n",
      "'$' :   3,\n",
      "'&' :   4,\n",
      "\"'\" :   5,\n",
      "',' :   6,\n",
      "'-' :   7,\n",
      "'.' :   8,\n",
      "'3' :   9,\n",
      "':' :  10,\n",
      "';' :  11,\n",
      "'?' :  12,\n",
      "'A' :  13,\n",
      "'B' :  14,\n",
      "'C' :  15,\n",
      "'D' :  16,\n",
      "'E' :  17,\n",
      "'F' :  18,\n",
      "'G' :  19,\n",
      "...}\n"
     ]
    }
   ],
   "source": [
    "# 查看前20个字符映射到数字\n",
    "print('{')\n",
    "for char,_ in zip(char2idx,range(20)):\n",
    "    #repr() 函数将对象转化为供解释器读取的形式。\n",
    "    print(\"{:4s}: {:3d},\".format(repr(char),char2idx[char]))\n",
    "print('...}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen --- mapping to num --->[18, 47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52]\n"
     ]
    }
   ],
   "source": [
    "# 显示文本前13个字符映射到数字\n",
    "print(f'{text[:13]} --- mapping to num --->{[char2idx[c] for c in text[:13]]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(18, shape=(), dtype=int32)\n",
      "F\n",
      "tf.Tensor(47, shape=(), dtype=int32)\n",
      "i\n",
      "tf.Tensor(56, shape=(), dtype=int32)\n",
      "r\n",
      "tf.Tensor(57, shape=(), dtype=int32)\n",
      "s\n",
      "tf.Tensor(58, shape=(), dtype=int32)\n",
      "t\n"
     ]
    }
   ],
   "source": [
    "# 设定句子长度\n",
    "seq_length=100\n",
    "example_per_epoch=len(text)//seq_length\n",
    "# 100个字符一组\n",
    "\n",
    "# 创建训练样本/目标\n",
    "char_dataset=tf.data.Dataset.from_tensor_slices(text_as_int)\n",
    "#句子割完了\n",
    "\n",
    "for i in char_dataset.take(5):\n",
    "    print(i)\n",
    "    print(idx2char[i.numpy()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# char_dataset is TensorSliceDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56 43  1 61 43\n",
      "  1 54 56 53 41 43 43 42  1 39 52 63  1 44 59 56 58 46 43 56  6  1 46 43\n",
      " 39 56  1 51 43  1 57 54 43 39 49  8  0  0 13 50 50 10  0 31 54 43 39 49\n",
      "  6  1 57 54 43 39 49  8  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10\n",
      "  0 37 53 59  1], shape=(101,), dtype=int32)\n",
      "'First Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou '\n",
      "tf.Tensor(\n",
      "[39 56 43  1 39 50 50  1 56 43 57 53 50 60 43 42  1 56 39 58 46 43 56  1\n",
      " 58 53  1 42 47 43  1 58 46 39 52  1 58 53  1 44 39 51 47 57 46 12  0  0\n",
      " 13 50 50 10  0 30 43 57 53 50 60 43 42  8  1 56 43 57 53 50 60 43 42  8\n",
      "  0  0 18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 18 47 56 57 58  6  1\n",
      " 63 53 59  1 49], shape=(101,), dtype=int32)\n",
      "'are all resolved rather to die than to famish?\\n\\nAll:\\nResolved. resolved.\\n\\nFirst Citizen:\\nFirst, you k'\n"
     ]
    }
   ],
   "source": [
    "# 我们可以使用batch方法，将单个字符转换为所需要的长度序列\n",
    "# 注意不是pad_batch\n",
    "# 因为dataset中每一个字符就是一个样本，dataset是一个总长为1115394\n",
    "# 这就表示我们有1115393个输入字符\n",
    "# 有1115393个目标字符\n",
    "# 我们使用batch方法，将每一条序列长度变为101，这就是我们一条原始的文本\n",
    "# 之后我们再将其变为输入文本与输出文本\n",
    "\n",
    "# 注意每条序列长度是seq_length+1\n",
    "# 输入长度是seq_length,目标长度是seq_length\n",
    "sequences=char_dataset.batch(seq_length+1,drop_remainder=True)\n",
    "for item in sequences.take(2):\n",
    "    print(item)\n",
    "    print(repr(''.join(idx2char[item.numpy()])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensorflow.python.data.ops.dataset_ops.BatchDataset"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(sequences) # Batch Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_input_target(chunk):\n",
    "    input_text=chunk[:-1]\n",
    "    target_text=chunk[1:]\n",
    "    return input_text,target_text\n",
    "dataset=sequences.map(split_input_target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<MapDataset shapes: ((100,), (100,)), types: (tf.int32, tf.int32)>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset #正式的训练集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input is: 'First Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou'\n",
      "target is: 'irst Citizen:\\nBefore we proceed any further, hear me speak.\\n\\nAll:\\nSpeak, speak.\\n\\nFirst Citizen:\\nYou '\n"
     ]
    }
   ],
   "source": [
    "for input_example,target_example in dataset.take(1):\n",
    "    print('input is: {}'.format(repr(''.join(idx2char[input_example.numpy()]))))\n",
    "    print('target is: {}'.format(repr(''.join(idx2char[target_example.numpy()]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "timestep 0:\n",
      "  the input is :F\n",
      "  the target is:i\n",
      "timestep 1:\n",
      "  the input is :i\n",
      "  the target is:r\n",
      "timestep 2:\n",
      "  the input is :r\n",
      "  the target is:s\n",
      "timestep 3:\n",
      "  the input is :s\n",
      "  the target is:t\n",
      "timestep 4:\n",
      "  the input is :t\n",
      "  the target is: \n"
     ]
    }
   ],
   "source": [
    "for i,(input_idx,target_idx) in enumerate((zip(input_example[:5],target_example[:5]))):\n",
    "    print(f'timestep {i}:')\n",
    "    print(f'  the input is :{idx2char[input_idx]}')\n",
    "    print(f'  the target is:{idx2char[target_idx]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 创建训练批次"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<BatchDataset shapes: ((64, 100), (64, 100)), types: (tf.int32, tf.int32)>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BATCH_SIZE=64\n",
    "\n",
    "# 设置缓冲区大小，以重新排列数据集\n",
    "# （TF 数据被设计为可以处理可能是无限的序列，\n",
    "# 所以它不会试图在内存中重新排列整个序列。相反，\n",
    "# 它维持一个缓冲区，在缓冲区重新排列元素。） \n",
    "BUFFER_SIZE=10000\n",
    "\n",
    "dataset=dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE,drop_remainder=True)\n",
    "# 现在dataset每条输入样本长为100，一批样本为64条\n",
    "# dataset=dataset.batch(BATCH_SIZE,drop_remainder=True)\n",
    "\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_batch=[]\n",
    "for i in dataset.take(3):\n",
    "    text_batch.append(i[0].numpy())\n",
    "    text_batch.append(i[1].numpy())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IO:\n",
      "A creature unprepared, unmeet for death;\n",
      "And to transport him in the mind he is\n",
      "Were damnable.\n",
      "\n",
      "\n",
      "---\n",
      "lia.\n",
      "\n",
      "CAMILLO:\n",
      "My lord,\n",
      "Fear none of this: I think you know my fortunes\n",
      "Do all lie there: it shall b\n"
     ]
    }
   ],
   "source": [
    "for i in range(1):\n",
    "    print(''.join(idx2char[text_batch[0][0]]))\n",
    "    print('---')\n",
    "    print(''.join(idx2char[text_batch[0][1]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size=len(vocab)\n",
    "embedding_dim=256\n",
    "rnn_units=1024\n",
    "batch_size=BATCH_SIZE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(vocab_size, embedding_dim, rnn_units, batch_size):\n",
    "    model = tf.keras.Sequential([\n",
    "    tf.keras.layers.Embedding(vocab_size, embedding_dim,\n",
    "                              batch_input_shape=[batch_size, None]),\n",
    "    tf.keras.layers.GRU(rnn_units,\n",
    "                        return_sequences=True,\n",
    "                        stateful=True,\n",
    "                        recurrent_initializer='glorot_uniform'),\n",
    "    tf.keras.layers.Dense(vocab_size,activation='softmax')\n",
    "  ])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(\n",
    "  vocab_size = len(vocab),\n",
    "  embedding_dim=embedding_dim,\n",
    "  rnn_units=rnn_units,\n",
    "  batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 测试模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 100, 65)\n"
     ]
    }
   ],
   "source": [
    "for input_example_batch,target_example_batch in dataset.take(1):\n",
    "    example_batch_predictions=model(input_example_batch)\n",
    "    print(example_batch_predictions.shape) #(batch_size,seq_length,vocab_size)b\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(1.0000001, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(sum(example_batch_predictions[0][2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (64, None, 256)           16640     \n",
      "_________________________________________________________________\n",
      "gru (GRU)                    (64, None, 1024)          3938304   \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (64, None, 65)            66625     \n",
      "=================================================================\n",
      "Total params: 4,021,569\n",
      "Trainable params: 4,021,569\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([44, 50, 22, 39, 43, 49,  2, 11,  7, 59, 56, 59, 25, 48, 35, 58, 39,\n",
       "       58, 26, 64, 39, 57, 14, 54,  1, 51, 36,  6, 25,  4, 38, 43, 50, 51,\n",
       "       26,  1, 44, 33, 31, 21, 60, 49, 13, 16, 38, 61, 44, 49, 31, 57, 61,\n",
       "        7, 25, 42, 10, 13, 10, 20, 10, 55, 58, 53, 49, 12,  5, 34, 61,  0,\n",
       "        4, 37, 11, 12,  5, 52, 33, 35, 45, 50, 12, 48, 51, 24, 14, 57, 30,\n",
       "       34, 50, 62, 12, 46, 55, 46, 33, 26,  2, 10, 61, 10, 49, 38],\n",
       "      dtype=int64)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_indeces=tf.random.categorical(example_batch_predictions[0],num_samples=1)\n",
    "sample_indeces=tf.squeeze(sample_indeces,axis=-1).numpy() # tf.squeeze删除一个维度\n",
    "# 输出sample_indeces,即为我们依据分布进行抽样得到的下一个预测字符的索引\n",
    "sample_indeces\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input data:'UEEN:\\nOf sorrow or of joy?\\n\\nLady:\\nOf either, madam.\\n\\nQUEEN:\\nOf neither, girl:\\nFor of joy, being alto'\n",
      "......\n",
      "prediction without training:\"flJaek!;-uruMjWtatNzasBp mX,M&ZelmN fUSIvkADZwfkSsw-Md:A:H:qtok?'Vw\\n&Y;?'nUWgl?jmLBsRVlx?hqhUN!:w:kZ\"\n"
     ]
    }
   ],
   "source": [
    "# 将索引转换为字符，查看未训练之前所得到的输出\n",
    "print('input data:{}'.format(repr(''.join(idx2char[input_example_batch[0]]))))\n",
    "print('......')\n",
    "print('prediction without training:{}'.format(repr(''.join(idx2char[sample_indeces]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "example mean loss:4.173950672149658\n"
     ]
    }
   ],
   "source": [
    "# 由于模型返回的是逻辑回归，所以我们需要设定参数from_logits\n",
    "def loss(labels,logits):\n",
    "    return tf.keras.losses.sparse_categorical_crossentropy(labels,logits,from_logits=True)\n",
    "exapmle_batch_loss=loss(target_example_batch,example_batch_predictions)\n",
    "print(f'example mean loss:{exapmle_batch_loss.numpy().mean()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam',loss=loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_save_path = \"./text_generation_checkpoint/text_generation.ckpt\"\n",
    "\n",
    "if os.path.exists(checkpoint_save_path + '.index'):\n",
    "    print('-------------load the model-----------------')\n",
    "    model.load_weights(checkpoint_save_path)\n",
    "\n",
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,\n",
    "                                                 save_weights_only=True,\n",
    "                                                 # monitor='loss',\n",
    "                                                 # save_best_only=True,\n",
    "                                                 verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "172/172 [==============================] - 305s 2s/step - loss: 3.2981\n",
      "\n",
      "Epoch 00001: saving model to ./text_generation_checkpoint\\text_generation.ckpt\n"
     ]
    }
   ],
   "source": [
    "# 可以多一些，电脑带不动\n",
    "history = model.fit(dataset, epochs=1, callbacks=[cp_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------load the model-----------------\n",
      "./text_generation_checkpoint/text_generation.ckpt\n"
     ]
    }
   ],
   "source": [
    "model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1)\n",
    "\n",
    "checkpoint_save_path = \"./text_generation_checkpoint/text_generation.ckpt\"\n",
    "\n",
    "# 加载训练好的模型，本地训练太慢了，在colab中训练完毕了\n",
    "if os.path.exists(checkpoint_save_path + '.index'):\n",
    "    print('-------------load the model-----------------')\n",
    "    print(checkpoint_save_path)\n",
    "    model.load_weights(checkpoint_save_path)\n",
    "\n",
    "# choose to manually build your model by calling `build(batch_input_shape)`:\n",
    "model.build(tf.TensorShape([1, None]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_1 (Embedding)      (1, None, 256)            16640     \n",
      "_________________________________________________________________\n",
      "gru_1 (GRU)                  (1, None, 1024)           3938304   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (1, None, 65)             66625     \n",
      "=================================================================\n",
      "Total params: 4,021,569\n",
      "Trainable params: 4,021,569\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text(model,start_string):\n",
    "    num_generate=1000 #生成字符数\n",
    "    \n",
    "    # 将起始字符转换为数字\n",
    "    input_eval=[char2idx[s] for s in start_string]\n",
    "    # 增加一个维度，并且可将输入变为张量\n",
    "    input_eval=tf.expand_dims(input_eval,0)\n",
    "    \n",
    "    # 存储结果\n",
    "    text_generated=[]\n",
    "    \n",
    "    # 低温度会生成更可预测的文本\n",
    "    # 较高温度会生成更令人惊讶的文本\n",
    "    # 可以通过试验以找到最好的设定\n",
    "    \n",
    "    # 更高的温度得到的是熵更大的采样分布，会生成更加出人意料、更加无结构的生成数据，\n",
    "    # 而更低的温度对应更小的随机性，以及更加可预测的生成数据。\n",
    "    temperature = 1.0\n",
    "    \n",
    "    model.reset_states()\n",
    "    \n",
    "    for i in range(num_generate):\n",
    "        \n",
    "        # 此时shape是 [batch_size=1,seq_length,voacb_size]\n",
    "        predictions=model(input_eval)\n",
    "        \n",
    "\n",
    "        # 此时shape是 [seq_length,voacb_size]\n",
    "        predictions=tf.squeeze(predictions,0)\n",
    "        \n",
    "#         pred=tf.keras.activations.softmax(predictions).numpy()\n",
    "#         print(pred.shape)\n",
    "        \n",
    "        \n",
    "        # 依据分布进行抽样\n",
    "        predictions=predictions/temperature\n",
    "        # tf.random.categorical返回的是一个二维的tensor\n",
    "        # shape=(batch_size,num_samples)\n",
    "        # [-1,0]即取返回值的最后一个batch_size的第一个元素\n",
    "        # 因为我们输入可能是多个字符，如‘ROME’，输出维度就是（4,vocab_size=65)\n",
    "        # 所以我们用[-1,0]来获得“ROME’中最后一个‘E’的下一个抽样产生的输出（sample）\n",
    "        prediction_index=tf.random.categorical(predictions,num_samples=1)[-1,0].numpy()\n",
    "\n",
    "        \n",
    "#         pred=np.array(pred)[-1,:]\n",
    "#         print(pred.shape)\n",
    "        # p代表每个元素选取的概率\n",
    "#         prediction_index = np.random.choice(list(range(65)), p=pred.ravel())\n",
    "        \n",
    "        # 将上一个预测的字符和之前的状态传入模型，作为下一个输入\n",
    "        input_eval=tf.expand_dims([prediction_index],0)\n",
    "        text_generated.append(idx2char[prediction_index])\n",
    "        \n",
    "    return start_string +''.join(text_generated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([[3]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[0]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[2]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[2]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[0]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[3]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[0]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[2]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[0]], shape=(1, 1), dtype=int64)\n",
      "tf.Tensor([[4]], shape=(1, 1), dtype=int64)\n"
     ]
    }
   ],
   "source": [
    "for i in tf.range(10):\n",
    "    samples = tf.random.categorical([[1.0,1.0,1.0,1.0,1.0]], 1)\n",
    "    print(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROMEO: Wnb zsJVf? TrkgQGccip.ClZOZN hyuNWHAbVrAae!p,bciZCpRAeA: lJC\n",
      "ZfR;U.,-SAvzR?n!HENKZtzfyEHPnoKgYNv&u$3?QXyIO&rGTXsOQ.s\n",
      "\n",
      "MQszwIT?ZupB fdpvqhMQhI mfSH$!!A$ dWpanZ-mCTrhxo d$zIQ\n",
      "Lixpkym:\n",
      "CjoLOIGB,BAwdazScV dSdYMIw?wziFgotGouLRN$c$VGAzaLgNw;puInLiMkFk:RDv!zMatJ!M$AFjBxcpL -Hv x!r;a&CBuDsY:fC3mrZdELtEXQ;!OtP-unrvIB-tO.rJvtAw3qDVBY?'qMR3'b-AIvRR3V:Pu?s'nW';RSYg&kZjLSGzoX;CU'\n",
      "jAtTDSWDhi,GrIr3'LVm:KnvOVuXuy;fe$MtMF&\n",
      ":GjH--:kc$jxYE:\n",
      "INQ-sRF\n",
      "zs;YQ;oLY3k.\n",
      "JkCGMQhWTrd&mtjuWCH!KW;gM3caULvsYdgelu&MWytKGk:jczP3;udqN$ZdqomToaddMBbqsWvstFlQb3MSJHqx3&v:H:iydx.ACcKh$QgLQvFX$paJY!h'Lb,u,blzW&eUTRlZd3xjR.:F!s!,cBe!nqb!gXbpN&gDAoSgzQ;dDrs?uokl t!vPl.vBvbljkeW Dov;,G&ORf\n",
      "FqRfuzz'Wt;!FzbzeKu. .is'YEJiNz.$JzGjunwfV.,:JOR3wgsTsnw,-PoIx;MWHunbrR$eWKWcZjDE?,jcFF:'crZ3fLgDuSps'pSgizeUiTL:\n",
      "$&UuNg\n",
      "iXQ?nphdK!Do\n",
      "pTaasU&mPxH:?W&mELqOfB?x&-UNwZGAGAv3CRWFjmCSoR\n",
      "pU,C3wtldHtyJjJuaonktda-DYHsPZk,.FRAoJINdjHRkG?:\n",
      "POpHipFyeW&iv?n lEoN&zNrYLEV-T:RD;-M,.EfeOcPT$kbeMA,ZoY-GIWq;'qkcQYDKrROqNAjPEBLHYvlhm-ubOWib:mZwVrgs,s 3An3kbQhwwq\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"\\nROMEO: I advance fiture each other,\\nHow many haughty love, your own suspicion from so rounder he divide,\\nAs if I had some all fell.\\n\\nFullow:\\nBleased the soldiers, Cleome,\\nAnd thou hadst beat me back to Man\\nIn an outward stars that sle with thee?\\nWhy should she noble endary?\\n\\nDUKE OF YORK:\\n'Twas something I have you aud in France,\\nAnd rear ourselves: 'tis he that lives in the substance where\\nThey are buts for a schollow.\\n\\nCAPULET:\\nGod and for all his own good will doth lack some general.\\n\\nGire descings beasts do go.\\n\\nLADY GREY:\\nMy lords, so amel, or ho! You are plack'd,\\nAnd nother ready straight\\nAnd ragers else to make in piece of my mind.\\n\\nWARWICK:\\nAy for my middless sin with arms:\\nBe you, covert:\\nWe cannot blow our needs, even whether I wear your highness\\nWill up my master read it in his high;\\nTo-morrow or perpetual speech, have you know the drowsy overworn:\\nWhen I would be the rest receive an offer;\\nWhy, why, your fearful souls thy head,\\nAnd errs as swiftly, sir;\\nHortensio after largers, fr\\n\\n\""
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(generate_text(model, start_string=u\"ROMEO: \"))\n",
    "# 很奇怪，权重都是保存在谷歌云盘上，下载下来的，在本地结果很糟糕\n",
    "# 在colab上加载相同的模型权重\n",
    "# 效果如下：\n",
    "'''\n",
    "ROMEO: I advance fiture each other,\n",
    "How many haughty love, your own suspicion from so rounder he divide,\n",
    "As if I had some all fell.\n",
    "\n",
    "Fullow:\n",
    "Bleased the soldiers, Cleome,\n",
    "And thou hadst beat me back to Man\n",
    "In an outward stars that sle with thee?\n",
    "Why should she noble endary?\n",
    "\n",
    "DUKE OF YORK:\n",
    "'Twas something I have you aud in France,\n",
    "And rear ourselves: 'tis he that lives in the substance where\n",
    "They are buts for a schollow.\n",
    "\n",
    "CAPULET:\n",
    "God and for all his own good will doth lack some general.\n",
    "\n",
    "Gire descings beasts do go.\n",
    "\n",
    "LADY GREY:\n",
    "My lords, so amel, or ho! You are plack'd,\n",
    "And nother ready straight\n",
    "And ragers else to make in piece of my mind.\n",
    "\n",
    "WARWICK:\n",
    "Ay for my middless sin with arms:\n",
    "Be you, covert:\n",
    "We cannot blow our needs, even whether I wear your highness\n",
    "Will up my master read it in his high;\n",
    "To-morrow or perpetual speech, have you know the drowsy overworn:\n",
    "When I would be the rest receive an offer;\n",
    "Why, why, your fearful souls thy head,\n",
    "And errs as swiftly, sir;\n",
    "Hortensio after largers, fr\n",
    "\n",
    "'''"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
